#include "precomp.h"

#ifdef __linux__

#include <sys/uio.h>
#include <stdio.h>
#include <stdlib.h>
#include <fstream>
#include <sstream>
#include <cstdint>

struct MemoryPermissions {
	uint8_t read : 1; // 读权限
	uint8_t write : 1; // 写权限
	uint8_t execute : 1; // 执行权限
	uint8_t shared : 1; // 共享标志
	uint8_t private_ : 1; // 私有标志
	uint8_t reserved : 3; // 保留位

	// 从字符串解析权限
	void parse(const std::string& perm_str)
	{
		read = (perm_str.find('r') != std::string::npos);
		write = (perm_str.find('w') != std::string::npos);
		execute = (perm_str.find('x') != std::string::npos);
		shared = (perm_str.find('s') != std::string::npos);
		private_ = (perm_str.find('p') != std::string::npos);
	}

	// 转换为可读字符串
	std::string to_string() const
	{
		std::string s;
		s += read ? 'r' : '-';
		s += write ? 'w' : '-';
		s += execute ? 'x' : '-';
		s += shared ? 's' : (private_ ? 'p' : '-');
		return s;
	}
};

struct MemoryRegion {
	uintptr_t start_addr;
	uintptr_t end_addr;
	MemoryPermissions permissions; // 替换原来的string
	off_t offset;
	std::string device;
	ino_t inode;
	std::string pathname;
	size_t size()
	{
		return end_addr - start_addr;
	}
};

std::vector<MemoryRegion> parse_proc_maps(pid_t pid)
{
	std::vector<MemoryRegion> regions;
	std::ifstream maps_file("/proc/" + std::to_string(pid) + "/maps");

	if (!maps_file.is_open()) {
		std::cerr << "Failed to open /proc/" << pid << "/maps" << std::endl;
		return regions;
	}

	std::string line;
	while (std::getline(maps_file, line)) {
		std::istringstream iss(line);
		MemoryRegion region;
		char dash;
		std::string addr_range, perm_str;

		// 解析地址范围
		iss >> addr_range;
		std::istringstream addr_ss(addr_range);
		addr_ss >> std::hex >> region.start_addr >> dash >> region.end_addr;

		// 解析权限
		iss >> perm_str;
		region.permissions.parse(perm_str);

		// 解析其他字段...
		iss >> std::hex >> region.offset;
		iss >> region.device;
		iss >> region.inode;
		std::getline(iss, region.pathname);
		region.pathname.erase(0, region.pathname.find_first_not_of(" \t"));

		regions.push_back(region);
	}

	return regions;
}

static int lua_openProcess(lua_State* L)
{
	int pid = lua_tointeger(L, 1);
	lua_pushinteger(L, (lua_Integer)pid);
	return 1;
}

static int lua_closeHandle(lua_State* L)
{
	return 0;
}

static int lua_getProcessMemMap(lua_State* L)
{
	pid_t pid = lua_tointeger(L, 1);
	std::vector<MemoryRegion> pmaps = parse_proc_maps(pid);
	lua_newtable(L);
	size_t idx = 1;
	for (MemoryRegion& m : pmaps) {
		lua_newtable(L);
		lua_pushliteral(L, "start");
		lua_pushinteger(L, m.start_addr);
		lua_rawset(L, -3);
		lua_pushliteral(L, "end");
		lua_pushinteger(L, m.end_addr);
		lua_rawset(L, -3);
		lua_pushliteral(L, "size");
		lua_pushinteger(L, m.size());
		lua_rawset(L, -3);
		lua_pushliteral(L, "device");
		lua_pushlstring(L, m.device.c_str(), m.device.size());
		lua_rawset(L, -3);
		lua_pushliteral(L, "path");
		lua_pushlstring(L, m.pathname.c_str(), m.pathname.size());
		lua_rawset(L, -3);
		std::string ps = m.permissions.to_string();
		lua_pushliteral(L, "permission");
		lua_pushlstring(L, ps.c_str(), ps.size());
		lua_rawset(L, -3);
		lua_rawseti(L, -2, idx);
		idx++;
	}
	return 1;
}

static int read_memory(pid_t pid, void* addr, size_t len, void* buf, size_t* got)
{
	struct iovec local[1];
	struct iovec remote[1];

	local[0].iov_base = buf;
	local[0].iov_len = len;
	remote[0].iov_base = addr;
	remote[0].iov_len = len;

	ssize_t nread = process_vm_readv(pid, local, 1, remote, 1, 0);
	if (nread == -1) {
		return -1;
	}
	*got = nread;
	return *got;
}

static int write_memory(pid_t pid, void* addr, size_t len, const void* buf, size_t* got)
{
	struct iovec local[1];
	struct iovec remote[1];

	local[0].iov_base = (void*)buf;
	local[0].iov_len = len;
	remote[0].iov_base = addr;
	remote[0].iov_len = len;

	ssize_t nwrite = process_vm_writev(pid, local, 1, remote, 1, 0);
	if (nwrite == -1) {
		return -1;
	}
	*got = nwrite;
	return *got;
}

CPointerList searchProcessString(pid_t pid, const char* searchString, size_t searchStringLength)
{
	CPointerList resultAddresses;
	if (searchStringLength <= 0)
		searchStringLength = strlen(searchString);
	std::vector<MemoryRegion> regions = parse_proc_maps(pid);
	char* buffer = 0;
	size_t buf_size = 0;
	for (MemoryRegion& memInfo : regions) {
		// 检查内存区域是否可读取
		if (!memInfo.permissions.read)
			continue;

		// 分配缓冲区来存储读取的内存
		if (buf_size < memInfo.size()) {
			buffer = (char*)realloc(buffer, memInfo.size());
			buf_size = memInfo.size();
			size_t bytesRead;

			// 读取进程内存
			if (read_memory(pid, (void*)memInfo.start_addr, memInfo.size(), buffer, &bytesRead) > 0) {
				// 在读取的内存中搜索字符串
				for (size_t i = 0; i < bytesRead - searchStringLength; i++) {
					if (memcmp(buffer + i, searchString, searchStringLength) == 0) {
						// 找到匹配，计算在进程中的实际地址
						void* matchAddress = (void*)((uintptr_t)memInfo.start_addr + i);
						resultAddresses.push_back(matchAddress);
					}
				}
			}
		}
	}
	if (buffer)
		free(buffer);
	return resultAddresses;
}

static int lua_readProcessMemory(lua_State* L)
{
	pid_t h = (pid_t)lua_tointeger(L, 1);
	void* start = (void*)lua_tointeger(L, 2);
	size_t size = lua_tointeger(L, 3);
	size_t got = 0;
	void* buffer = malloc(size);
	if (!buffer)
		return 0;
	if (read_memory(h, start, size, buffer, &got) < 0) {
		free(buffer);
		return 0;
	}
	lua_pushlstring(L, (char*)buffer, got);
	free(buffer);
	return 1;
}

static int lua_writeProcessMemory(lua_State* L)
{
	pid_t h = (pid_t)lua_tointeger(L, 1);
	void* start = (void*)lua_tointeger(L, 2);
	size_t size;
	const char* buffer = luaL_checklstring(L, 3, &size);
	size_t got = 0;
	if (write_memory(h, start, size, buffer, &got) < 0)
		return 0;
	lua_pushinteger(L, got);
	return 1;
}

/*
add memory operation funcs
module table is on top of stack on entry
*/
static int add_memory_ops(lua_State* L)
{
	// enumProcesses is in linux.lua
	// regfunc(enumProcesses);
	regfunc(getProcessMemMap);
	regfunc(openProcess);
	regfunc(closeHandle);
	regfunc(readProcessMemory);
	regfunc(writeProcessMemory);
	int target = lua_gettop(L);
	luartk_embed_script_with_ret(L, "linux");
	luartk_merge_table(L, target);
	return 0;
}

GAMEX_ADD_INIT_FUNC(add_memory_ops)

#endif // __linux__
